"""
Barotropic Model for use with the Variational Data Assimilation Framework
Implements a two-dimensional barotropic vorticity equation model
"""
import numpy as np
from scipy.ndimage import gaussian_filter, map_coordinates
import time

class BarotropicModel:
    """
    A simple barotropic vorticity equation model for atmospheric flow.
    This model represents non-divergent 2D flow on a beta-plane.
    """
    def __init__(self, nx=3, ny=3, dx=100000, dy=100000, dt=300, 
                 beta=1e-12, diffusion=10):
        """
        Initialize barotropic model
        
        Args:
            nx, ny: Grid dimensions
            dx, dy: Grid spacing in meters
            dt: Time step in seconds
            beta: Beta-plane parameter (df/dy)
            diffusion: Diffusion coefficient
        """
        self.nx = nx
        self.ny = ny
        self.state_size = nx * ny
        self.dx = dx
        self.dy = dy
        self.dt = dt
        self.beta = beta
        self.diffusion = diffusion
        
        # Initialize model parameters and grids
        self._setup_grid()
        # Precompute grid indices for interpolation
        self.grid_x = np.arange(nx)
        self.grid_y = np.arange(ny)
        
        # Initialize timing dictionary
        self.timings = {}
        
    def _setup_grid(self):
        """Setup spatial grids for the model"""
        self.x = np.linspace(0, self.nx * self.dx, self.nx)
        self.y = np.linspace(0, self.ny * self.dy, self.ny)
        self.X, self.Y = np.meshgrid(self.x, self.y)
        
        # Precompute matrices for Poisson solver (streamfunction from vorticity)
        kx = np.fft.fftfreq(self.nx, d=self.dx/(2*np.pi))
        ky = np.fft.fftfreq(self.ny, d=self.dy/(2*np.pi))
        self.KX, self.KY = np.meshgrid(kx, ky)
        self.K2 = self.KX**2 + self.KY**2
        # Avoid division by zero
        self.K2[0, 0] = 1.0
        
    def _reshape_state(self, x_flat):
        """Reshape flat state vector to 2D grid"""
        return x_flat.reshape(self.ny, self.nx)
        
    def _flatten_state(self, x_2d):
        """Flatten 2D grid to state vector"""
        return x_2d.flatten()
    
    def _compute_streamfunction(self, vorticity):
        """Solve Poisson equation ∇²ψ = ζ for streamfunction ψ"""
        vort_fft = np.fft.fft2(vorticity)
        psi_fft = -vort_fft / self.K2
        psi_fft[0, 0] = 0  # Remove mean
        return np.real(np.fft.ifft2(psi_fft))
    
    def _compute_velocity(self, psi):
        """Compute velocity components from streamfunction"""
        # u = -∂ψ/∂y, v = ∂ψ/∂x
        u = -np.gradient(psi, self.dy, axis=0)
        v = np.gradient(psi, self.dx, axis=1)
        return u, v
    
    def _apply_diffusion(self, vorticity):
        """Apply diffusion to vorticity field"""
        sigma = np.sqrt(2 * self.diffusion * self.dt) / self.dx
        return gaussian_filter(vorticity, sigma=sigma)
    
    def step(self, x):
        """Step forward in time using semi-Lagrangian scheme"""
        start_time = time.time()
        
        # Reshape state vector to 2D grid
        vorticity = self._reshape_state(x)
        
        # Calculate streamfunction from vorticity
        t0 = time.time()
        psi = self._compute_streamfunction(vorticity)
        self.timings['compute_streamfunction'] = self.timings.get('compute_streamfunction', 0) + (time.time() - t0)
        
        # Calculate velocity field
        t0 = time.time()
        u, v = self._compute_velocity(psi)
        self.timings['compute_velocity'] = self.timings.get('compute_velocity', 0) + (time.time() - t0)
              
        # Vectorized departure point calculation and interpolation
        t0 = time.time()
        Xd = (self.X - u*self.dt) % (self.nx*self.dx)
        Yd = (self.Y - v*self.dt) % (self.ny*self.dy)
        
        # Convert to grid indices
        xi = (Xd / self.dx).astype(np.float32)
        yi = (Yd / self.dy).astype(np.float32)
        
        # Vectorized interpolation using precomputed grid
        vorticity_new = map_coordinates(
            vorticity, 
            [yi.ravel(), xi.ravel()],
            order=3,
            mode='wrap'
        ).reshape(vorticity.shape)
        self.timings['interpolation'] = self.timings.get('interpolation', 0) + (time.time() - t0)
        
        # Vectorized physics updates (beta effect)
        vorticity_new += self.beta * v * self.dt
        
        # Apply diffusion
        t0 = time.time()
        result = self._flatten_state(self._apply_diffusion(vorticity_new))
        self.timings['diffusion'] = self.timings.get('diffusion', 0) + (time.time() - t0)
        
        # Track total step time
        self.timings['step_total'] = self.timings.get('step_total', 0) + (time.time() - start_time)
        
        return result

    def generate_random_state(self, mean_amplitude=1.0, seed=None):
        if seed is not None:
            np.random.seed(seed)
            
        kmax = min(self.nx, self.ny) // 2
        kx = np.arange(-kmax, kmax+1)
        ky = np.arange(-kmax, kmax+1)
        kx_grid, ky_grid = np.meshgrid(kx, ky, indexing='ij')
        
        # Vectorized spectral calculations
        k = np.sqrt(kx_grid**2 + ky_grid**2)
        mask = (k > 1) & (k <= kmax)
        
        spectrum = np.zeros((self.ny, self.nx), dtype=np.complex128)
        phs = 2j * np.pi * np.random.rand(*kx_grid.shape)
        epsilon = 1e-10
        amp = np.where((mask) & (k > 0), mean_amplitude * (k + epsilon)**-1.5, 0)
        
        # Vectorized index mapping
        idx_x = kx_grid % self.nx
        idx_y = ky_grid % self.ny
        spectrum[idx_y, idx_x] = amp * np.exp(phs)

        return self._flatten_state(np.real(np.fft.ifft2(spectrum)))
        

    def integrate(self, x0, t_steps):
        """Integrate model forward for t_steps"""
        start_time = time.time()
        
        trajectory = np.zeros((t_steps + 1, len(x0)))
        trajectory[0] = x0
        
        for t in range(t_steps):
            trajectory[t + 1] = self.step(trajectory[t])
        
        self.timings['integrate'] = self.timings.get('integrate', 0) + (time.time() - start_time)
        return trajectory
    def tangent_linear(self, x_ref, epsilon=1e-20):
        """Compute Jacobian matrix using complex step"""
        n = len(x_ref)
        J = np.zeros((n,n))
        for i in range(n):
            e_i = np.zeros(n); e_i[i] = 1
            perturbed_state = x_ref + e_i * (epsilon*1j)
            J[:,i] = self.step(perturbed_state).imag / epsilon
        return J
    
    def adjoint(self, x_ref, epsilon=1e-20):
        """Adjoint Jacobian via transpose"""
        return self.tangent_linear(x_ref, epsilon).conj().T

    def validate_adjoint(self, x_ref, tol=1e-6):
        """Validate adjoint implementation using dot product test"""
        v = np.random.randn(self.state_size)
        w = np.random.randn(self.state_size)
        
        # Compute Lv · w
        L = self.tangent_linear(x_ref)
        lhs = np.dot(L @ v, w)
        
        # Compute v · L*w
        L_adj = self.adjoint(x_ref)
        rhs = np.dot(v, L_adj @ w)
        
        # Calculate relative error
        rel_error = abs(lhs - rhs) / (abs(lhs) + abs(rhs) + 1e-20)
        
        if rel_error < tol:
            return True, rel_error
        return False, rel_error
       
    def get_timings(self):
        """Return the timing statistics for the model functions"""
        return self.timings
    def print_timing_statistics(self):
        """Print detailed statistics about function execution times"""
        if not self.timings:
            print("No timing data available.")
            return
            
        # Calculate total model execution time excluding the integrate time
        operation_times = {k: v for k, v in self.timings.items() if k != 'integrate'}
        total_time = sum(operation_times.values())
        
        print("\n===== Timing Statistics =====")
        print(f"Total execution time: {total_time:.6f} seconds")
        
        # Sort operations by time (descending)
        sorted_ops = sorted(operation_times.items(), key=lambda x: x[1], reverse=True)
        
        print("\nBreakdown by operation:")
        print("-----------------------")
        for op, time_spent in sorted_ops:
            percentage = (time_spent / total_time) * 100
            print(f"{op:<25}: {time_spent:.6f} s ({percentage:.2f}%)")
